import os
import cv2
import re
import numpy as np


def downsample_image(image, scale=0.125):
    # 下采样图片
    downsampled_image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4)
    return downsampled_image


def batch_downsample_and_merge(input_dir, output_file, tile_size=1024, downscale=0.125):
    image_files = [f for f in os.listdir(input_dir) if f.endswith('.png')]

    # 使用正则表达式提取行列信息
    pattern = re.compile(r'S\d+M\d+-\d+_tr(\d+)-tc(\d+)\.png')
    row_col_info = [pattern.findall(f)[0] for f in image_files]
    rows = sorted(list(set(int(info[0]) for info in row_col_info)))
    cols = sorted(list(set(int(info[1]) for info in row_col_info)))

    # 计算大图尺寸
    num_rows = len(rows)
    num_cols = len(cols)
    downsampled_tile_size = int(tile_size * downscale)
    large_image = np.zeros((num_rows * downsampled_tile_size, num_cols * downsampled_tile_size, 3), dtype=np.uint8)

    for image_file in image_files:
        input_path = os.path.join(input_dir, image_file)
        image = cv2.imread(input_path)
        if image is None:
            print(f"无法打开图像文件：{input_path}")
            continue

        # 下采样图片
        downsampled_image = downsample_image(image, downscale)

        # 提取行列信息
        row_col = pattern.findall(image_file)[0]
        row = int(row_col[0]) - 1
        col = int(row_col[1]) - 1

        # 将下采样后的图片放置到大图的相应位置
        start_row = row * downsampled_tile_size
        start_col = col * downsampled_tile_size
        large_image[start_row:start_row + downsampled_tile_size,
        start_col:start_col + downsampled_tile_size] = downsampled_image

    # 保存合并后的大图
    cv2.imwrite(output_file, large_image)
    print(f"Processed large image saved to {output_file}")


# 示例用法
input_dir = 'E:\\wafer52\\11867_32nm'  # 替换为实际输入目录路径
output_file = 'E:\\wafer52\\11867_256nm\\S11866M3-163_tr1-tc1.png'  # 替换为实际输出路径

batch_downsample_and_merge(input_dir, output_file, tile_size=1024, downscale=0.125)

print("Batch downsampling and merging completed.")
